package com.example.mcp.client;

import com.example.mcp.server.model.McpInitializeRequest;
import com.example.mcp.server.model.McpToolCallRequest;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;
import org.springframework.web.socket.handler.TextWebSocketHandler;

public class McpWebSocketClient extends TextWebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(McpWebSocketClient.class);

    private final String wsUrl;
    private final ObjectMapper objectMapper = new ObjectMapper();
    private WebSocketSession session;

    public McpWebSocketClient(String wsUrl) {
        this.wsUrl = wsUrl;
    }

    public void connect() throws Exception {
        StandardWebSocketClient client = new StandardWebSocketClient();
        CompletableFuture<WebSocketSession> future = client.execute(this, wsUrl);
        this.session = future.get();
        logger.info("Connected to MCP WebSocket server at {}", wsUrl);
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) {
        try {
            String payload = message.getPayload();
            logger.debug("Received message: {}", payload);
        } catch (Exception e) {
            logger.error("Error handling received message: {}", e.getMessage(), e);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) throws Exception {
        logger.info("WebSocket connection closed with status: {}", status);
        this.session = null;
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        logger.error("WebSocket transport error: {}", exception.getMessage(), exception);
        this.session = null;
    }

    public void sendInitialize() throws Exception {
        if (session == null || !session.isOpen()) {
            throw new IllegalStateException("WebSocket session is not open");
        }

        McpInitializeRequest request = new McpInitializeRequest();
        request.setId(UUID.randomUUID().toString());
        request.setProtocolVersion("2024-11-05");
        request.setType("initialize");

        McpInitializeRequest.ClientCapabilities capabilities = new McpInitializeRequest.ClientCapabilities();
        capabilities.setExperimental(null);
        capabilities.setSampling(null);
        request.setCapabilities(capabilities);

        McpInitializeRequest.ClientInfo clientInfo = new McpInitializeRequest.ClientInfo();
        clientInfo.setName("Spring AI MCP WebSocket Client");
        clientInfo.setVersion("1.0.0");
        request.setClientInfo(clientInfo);

        String json = objectMapper.writeValueAsString(request);
        sendMessage(json);
    }

    public void callTool(String toolName, Map<String, Object> arguments) throws Exception {
        if (session == null || !session.isOpen()) {
            throw new IllegalStateException("WebSocket session is not open");
        }

        McpToolCallRequest request = new McpToolCallRequest();
        request.setId(UUID.randomUUID().toString());
        request.setType(toolName);
        request.setName(toolName);
        request.setArguments(arguments);

        String json = objectMapper.writeValueAsString(request);
        sendMessage(json);
    }

    private void sendMessage(String message) throws Exception {
        try {
            if (session != null && session.isOpen()) {
                session.sendMessage(new TextMessage(message));
                logger.debug("Sent message: {}", message);
            } else {
                logger.warn("Cannot send message, WebSocket session is not open");
                throw new IllegalStateException("WebSocket session is not open");
            }
        } catch (Exception e) {
            logger.error("Error sending message: {}", e.getMessage(), e);
            if (isConnectionReset(e)) {
                logger.info("Connection reset detected, clearing session");
                this.session = null;
            }
            throw e;
        }
    }

    private boolean isConnectionReset(Exception e) {
        return e instanceof java.io.IOException &&
                (e.getMessage().contains("Connection reset by peer") ||
                        e.getMessage().contains("Broken pipe"));
    }

    public void disconnect() throws Exception {
        if (session != null && session.isOpen()) {
            session.close();
            logger.info("Disconnected from MCP WebSocket server");
        } else {
            logger.info("WebSocket session was already closed or null");
        }
        this.session = null;
    }

    public boolean isConnected() {
        return session != null && session.isOpen();
    }
}